#include "stream.h"

static fftw_plan	fp1, bp1, fp2, bp2, ffb;
static double		*kx, *ky, *x, *y, visc;
static double		dt, lx, ly, k2;
static int		nx, ny;
static int		j, k, timeSeed;
static vect		aux, vel;
static fftw_complex	*tmp, *forcing, *rk[4], *v0, *lap;

void init_fft(ctrl_ptr in)
{
  double cx, cy;

  visc = in->visc; 
  nx = in->nx; ny = in->ny;
  lx = in->lx; ly = in->ly;
  cx = 2*pi/lx; cy = 2*pi/ly; 
  if ((ly/ny) > (lx/nx)) dt = (in->cfl)*(lx/nx);
  else dt = (in->cfl)*(ly/ny);


  aux.u = malloc(sizeof(fftw_complex)*nx*ny);
  aux.v = malloc(sizeof(fftw_complex)*nx*ny);
  vel.u = malloc(sizeof(fftw_complex)*nx*ny);
  vel.v = malloc(sizeof(fftw_complex)*nx*ny);
  tmp = malloc(sizeof(fftw_complex)*nx*ny);
  lap = malloc(sizeof(fftw_complex)*nx*ny);
  v0 = malloc(sizeof(fftw_complex)*nx*ny);
  forcing = malloc(sizeof(fftw_complex)*nx*ny);
  for (j = 0; j < 4; j++ ) {
    rk[j] = malloc(sizeof(fftw_complex)*nx*ny);
  }
  kx  = malloc(sizeof(fftw_complex)*nx);
  ky  = malloc(sizeof(fftw_complex)*ny);
  x  = malloc(sizeof(fftw_complex)*nx);
  y  = malloc(sizeof(fftw_complex)*ny);
  
  fp1 = fftw_plan_dft_2d(nx, ny, aux.u, aux.u, FFTW_FORWARD, FMODE);
  bp1 = fftw_plan_dft_2d(nx, ny, aux.u, aux.u, FFTW_BACKWARD, FMODE);
  fp2 = fftw_plan_dft_2d(nx, ny, aux.v, aux.v, FFTW_FORWARD, FMODE);
  bp2 = fftw_plan_dft_2d(nx, ny, aux.v, aux.v, FFTW_BACKWARD, FMODE);   
  ffb = fftw_plan_dft_2d(nx, ny, forcing, forcing, FFTW_BACKWARD, FMODE);   

  timeSeed = (int) time(0);
  srand(timeSeed);
  //printf("Debug:\t Setting seed to %d\nDebug:\t Random Number %e\n", timeSeed, 1.*rand()/RAND_MAX);
  for (j = 0; j < nx; j++) {
    if (j > nx/2) kx[j] = (j-nx)*cx;
    else kx[j] = j*cx;
  }
  for (j = 0; j < ny; j++) {
    if (j > ny/2) ky[j] = (j-ny)*cy;
    else ky[j] = j*cy;
  }
  for (j = 0; j < nx; j++) {
    x[j] = -0.5*lx + (j-0.5)*(lx/nx);
    for (k = 0; k < ny; k++){
      y[k] = -0.5*ly + (k-0.5)*(ly/ny);
      k2 = sqrt(kx[j]*kx[j] + ky[k]*ky[k]);
      forcing[k + ny*j] = (in->a_force)*(k2/(in->k_force))*cexp( -   pow(k2/(in->k_force),2)  );
    }
  }
  /*FILE *fhdlog = fopen("debug_forcing_0.txt","w");
  for (j = 0; j < nx; j++) {
    for (k = 0; k < ny; k ++) {     
      fprintf(fhdlog,"%e\t%e\t%e\t%e\n",kx[j], ky[k], cabs(forcing[k + ny*j]), atan2(cimag(aux.u[k + ny*j]), creal(aux.u[k + ny*j]) ));
    }
    fprintf(fhdlog,"\n");
  }
  fclose(fhdlog);*/
  fftw_execute(ffb);
}

void compute_velocity(ctrl_ptr in)
{
  for (j = 0; j < nx*ny; j++){ 
    aux.u[j] = in->vrt[j];
  }
  fftw_execute(fp1);
  for (j = 0; j < nx; j++){
    for (k = 0; k < ny; k++){
      tmp[k + ny*j] = aux.u[k + ny*j];
      k2 = kx[j]*kx[j] + ky[k]*ky[k];
      aux.u[k + ny*j] = 1.I*kx[j]*tmp[k + ny*j]/(k2*nx*ny);
      aux.v[k + ny*j] = 1.I*ky[k]*tmp[k + ny*j]/(k2*nx*ny); 
    }
  }
  aux.u[0] = 0;
  aux.v[0] = 0;
  fftw_execute(bp1);
  fftw_execute(bp2);
  for (j = 0; j < nx*ny; j++){
      in->vel.u[j] = aux.v[j];	
      in->vel.v[j] = -aux.u[j];
  }
}

void compute_gradv(ctrl_ptr in)
{
  for (j = 0; j < nx*ny; j++){ 
    aux.u[j] = in->vrt[j];
  }
  fftw_execute(fp1);
  for (j = 0; j < nx; j++){
    for (k = 0; k < ny; k++){
      tmp[k + ny*j] = aux.u[k + ny*j];
      aux.u[k + ny*j] = 1.I*kx[j]*tmp[k + ny*j]/(nx*ny);
      aux.v[k + ny*j] = 1.I*ky[k]*tmp[k + ny*j]/(nx*ny);
    }
  }
  aux.u[0] = 0;
  aux.v[0] = 0;
  fftw_execute(bp1);
  fftw_execute(bp2);
  for (j = 0; j < nx*ny; j++){
      (in->grad_v.u[j]) = aux.u[j];
      (in->grad_v.v[j]) = aux.v[j];		
  }
}
void compute_lapv(ctrl_ptr in)
{
  for (j = 0; j < nx*ny; j++) aux.u[j] = in->vrt[j]; 
  fftw_execute(fp1);
  for (j = 0; j < nx; j++)
    for (k = 0; k < ny; k ++) {     
       k2 = pow(kx[j]*kx[j] + ky[k]*ky[k], in->p);
       aux.u[k + ny*j] = -k2*aux.u[k + ny*j]/(nx*ny);
    }
  fftw_execute(bp1);
  for (j = 0; j < nx*ny; j++) lap[j] = aux.u[j]; 
}

void rhs(ctrl_ptr in, fftw_complex *out)
{
  compute_velocity(in);
  compute_gradv(in);
  for (j = 0; j < nx*ny; j++){ 
    out[j] = -( (in->vel.u[j])*(in->grad_v.u[j]) + (in->vel.v[j])*(in->grad_v.v[j]) ) + (in->FFlag)*forcing[j];
  }
}
void rhs2(ctrl_ptr in, fftw_complex *out)
{
  compute_velocity(in);
  compute_gradv(in);
  compute_lapv(in);
  for (j = 0; j < nx*ny; j++){ 
    out[j] = -( (in->vel.u[j])*(in->grad_v.u[j]) + (in->vel.v[j])*(in->grad_v.v[j]) ) + visc*lap[j] + 0.*(in->FFlag)*forcing[j];
  }
}

/*void dealias(ctrl_ptr in)
{
  for (j = 0; j < nx*ny; j++) aux.u[j] = in->vrt[j];
  fftw_execute(fp1);
  for (j = 0; j < nx; j++){
    for (k = 0; k < ny; k++){
      aux.u[k + ny*j] = aux.u[k + ny*j]*dealias[k + ny*j];
    }
  }
  
}*/
//--------------------------------------------
void rk4_nl(ctrl_ptr in, double mul){
  
  rhs(in, rk[0]);
  for (j = 0; j < nx*ny; j++) {
    v0[j] = in->vrt[j]; 
    in->vrt[j] = v0[j] + 0.5*mul*dt*rk[0][j];
  }
  rhs(in, rk[1]);
  for (j = 0; j < nx*ny; j++) in->vrt[j] = v0[j] + 0.5*mul*dt*rk[1][j]; 
  rhs(in, rk[2]);
  for (j = 0; j < nx*ny; j++) in->vrt[j] = v0[j] + dt*mul*rk[2][j]; 
  rhs(in, rk[3]);
  for (j = 0; j < nx*ny; j++) in->vrt[j] = v0[j] + (dt*mul/6.)*(rk[0][j]+2.*rk[1][j]+2.*rk[2][j]+rk[3][j]);
}

void rk4(ctrl_ptr in, double mul){
  
  for (j = 0; j < nx*ny; j++) aux.u[j] = forcing[j];
  fftw_execute(fp1);
  for (j = 0; j < nx*ny; j++) { 
    aux.u[j] = aux.u[j]*cexp(2.I*pi*(double)rand()/RAND_MAX)/(nx*ny);
  }
  fftw_execute(bp1);
  /*FILE *fhdlog = fopen("debug_forcing.txt","w");
  for (j = 0; j < nx; j++) {
    for (k = 0; k < ny; k ++) {     
      fprintf(fhdlog,"%e\t%e\t%e\n",x[j], y[k], cabs(aux.u[k + ny*j]));
    }
    fprintf(fhdlog,"\n");
  }
  fclose(fhdlog);*/
  for (j = 0; j < nx*ny; j++) forcing[j] = aux.u[j];




  rhs2(in, rk[0]);
  for (j = 0; j < nx*ny; j++) {
    v0[j] = in->vrt[j]; 
    in->vrt[j] = v0[j] + 0.5*mul*dt*(rk[0][j] + (in->FFlag)*forcing[j]);
  }
  rhs2(in, rk[1]);
  for (j = 0; j < nx*ny; j++) in->vrt[j] = v0[j] + 0.5*mul*dt*(rk[1][j] + (in->FFlag)*forcing[j]  ); 
  rhs2(in, rk[2]);
  for (j = 0; j < nx*ny; j++) in->vrt[j] = v0[j] + dt*mul*(rk[2][j] + (in->FFlag)*forcing[j]  ); 
  rhs2(in, rk[3]);
  for (j = 0; j < nx*ny; j++) in->vrt[j] = v0[j] + (dt*mul/6.)*(rk[0][j]+2.*rk[1][j]+2.*rk[2][j]+rk[3][j] + 6.*(in->FFlag)*forcing[j]);
}



void diffusion_step(ctrl_ptr in, double mul)
{
  for (j = 0; j < nx*ny; j++) aux.u[j] = in->vrt[j]; 
  fftw_execute(fp1);
  for (j = 0; j < nx; j++)
    for (k = 0; k < ny; k ++) {     
       k2 = kx[j]*kx[j] + ky[k]*ky[k];
       //in->sp[k + ny*j] = aux.u[k + ny*j]/(nx*ny);
       aux.u[k + ny*j] = aux.u[k + ny*j]*cexp(-mul*visc*k2*dt)/(nx*ny);
    }
  fftw_execute(bp1);
  for (j = 0; j < nx*ny; j++) in->vrt[j] = aux.u[j]; 
}

void find_spectrum(ctrl_ptr in)
{
  for (j = 0; j < nx*ny; j++) aux.u[j] = in->vrt[j]; 
  fftw_execute(fp1);
  for (j = 0; j < nx*ny; j++) in->sp[j] = aux.u[j]/(nx*ny);
  
}


//--------------------------------------------

void generate_ic(ctrl_ptr in)
{
  double a = in->a;
  double r2;
  for (j = 0; j < nx; j++){
    for (k = 0; k < ny; k++){
      //r2 = pow(-0.5*lx + (j-0.5)*(lx/nx) , 2) + pow(-0.5*ly + (k-0.5)*(ly/ny) , 2);
      in->vrt[k + ny*j] = 0; //(in->U)/a*(2. - r2/(a*a))*cexp(0.5*(1 - r2/(a*a) ));
    }
  }
}







